/////////////////////////
// File: camera_array_calibration.cpp
// Authors: Brian Fehrman and Scott Logan and Matt Richard
// Date: 3-27-2013
// 
// Purpose: Implementation file for the camera_array_calibartion functions
//       which are used for calibrating the cameras
//
/////////////////////////

////////// Includes //////////////
#include "camera_array_calibration.h"

/////////////////////////
// function: capture_images
// Authors: Brian Fehrman and Scott Logan and Matt Richard
// 
// Purpose: Captures images from the 2x2 array in roughly 5
//       second intervals until the user hits a key. Saves
//       The images named with row_col number, where the perspective
//       is if the user is behind the camera array and seeing
//       the world from the perspective of the array.
//       A video feed is shown for the user.
//
//
//////////////////////////
int capture_images( FileNode fs )
{  
   string base_dir;
   int delay_ms;
   int width, height;
   Mat cam_layout;
   vector< int > vid_list;
   
   fs["base_dir"] >> base_dir;
   fs["delay_ms"] >> delay_ms;
   fs["width"] >> width;
   fs["height"] >> height;
   fs["cam_layout"] >> cam_layout;
   
   clock_t clk_start;
   float elapsed = 0.0;
   float delay_s = delay_ms / 1000.0;
   int num_cams = cam_layout.rows * cam_layout.cols;
   int counter = 0;
   int num_pic = 0;
   stringstream pix_res;
   VideoCapture cam[ num_cams ];
   vector<string> folder_list;
   string combined_dir;
   
   mkdir( base_dir.c_str(), S_IRWXU|S_IRGRP|S_IXGRP );
   base_dir += "/";
   pix_res << width << "_" << height << ".jpg";
   
   combined_dir = base_dir + "combined";
   mkdir( combined_dir.c_str(), S_IRWXU|S_IRGRP|S_IXGRP );
   combined_dir += "/";
   
   //Flatten matrix for easier traversing and create
   //folder for each camera. -1 indicates no cam in that
   //spot
   for( int row = 0; row < cam_layout.rows; row++ )
   {
      for( int col = 0; col < cam_layout.cols; col++ )
      {
         //Don't create folder for cameras designated with -1
         if( cam_layout.at< int >( row, col ) > 0  )
         {
            stringstream cam_folder;
            cam_folder << base_dir <<  "cam_" << row << "_" << col;
            mkdir( cam_folder.str().c_str(), S_IRWXU|S_IRGRP|S_IXGRP );
            
            folder_list.push_back( cam_folder.str() + "/" ); 
         }
         else
         {
            //Add empty string so indexing is easier later
            folder_list.push_back("");
         }
         
         vid_list.push_back( cam_layout.at< int >( row, col ) );
      }
   }
   
   //Open the video feeds and set the resolution, ignore cameras
   //designated with -1
   for( int vid = 0; vid < num_cams; vid++ )
   {
      if( vid_list[ vid ]  > 0 )
      {
         cam[ vid ].open( vid_list[ vid ] );
         
         cam[ vid ].set( CV_CAP_PROP_FRAME_WIDTH, width);
         cam[ vid ].set( CV_CAP_PROP_FRAME_HEIGHT, height );
      }
   }

   namedWindow( "combined_gray", CV_WINDOW_NORMAL | CV_WINDOW_KEEPRATIO );
   namedWindow( "combined_color", CV_WINDOW_NORMAL | CV_WINDOW_KEEPRATIO );
   
   clk_start = clock();       
   
   //Capture loop. Goes until user hits a key.
   for( ;; )
   {
      Mat cap_color[ num_cams ];
      Mat cap_gray[ num_cams ];
      Mat combined_gray( height * cam_layout.rows, width * cam_layout.cols, CV_8UC1 );
      Mat combined_color( height * cam_layout.rows, width * cam_layout.cols, CV_8UC3); 
   
      //Capture camera images in parralel and convert to grayscale
      #pragma omp parallel for
      for( int vid = 0; vid < num_cams; vid++)
      {
         cap_gray[ vid ].create( height, width, CV_8UC1 );
         cap_color[ vid ].create( height, width, CV_8UC3 );
         
         if( vid_list[ vid ] > 0 )
         {
            cam[ vid ] >> cap_color[ vid ];
            cvtColor( cap_color[ vid ], cap_gray[ vid ], CV_BGR2GRAY );
         }
      }
      //For each camera in the array
      for( int row_comb = 0; row_comb < cam_layout.rows; row_comb++ )
      {
         for( int col_comb = 0; col_comb < cam_layout.cols; col_comb++ )
         {                   
            //Combine the image feeds in to single displays for the user
            for(int row = 0; row < height; row++ )
            {
               for( int col = 0; col < width; col++ )
               {
                  combined_gray.at<unsigned char>( row + ( height * row_comb ), 
                                          col + ( width * col_comb ) )
                                          = cap_gray[ col_comb + ( row_comb * cam_layout.cols ) ].at<unsigned char>( row, col ); 

                  combined_color.at<Vec3b>( row + ( height * row_comb ), 
                                    col + ( width * col_comb ) )
                                    = cap_color[ col_comb + ( row_comb * cam_layout.cols ) ].at<Vec3b>( row, col );
               }
            }
         }
      }
            
      imshow( "combined_color", combined_color );
      imshow( "combined_gray", combined_gray );
      
      elapsed = ( float( clock() - clk_start ) / CLOCKS_PER_SEC );
      
      //If desired delay time has passed then save the current set of images
      if( elapsed >= delay_s )
      {
         clk_start = clock();
         stringstream ss;
         
         ss << num_pic;
         cout << endl << endl;
         for( int vid = 0; vid < num_cams; vid++ )
         {
            if( vid_list[ vid ]  > 0 )
            {
               string color_save_string = folder_list[ vid ] + ss.str() + "_color_" + pix_res.str();
               string gray_save_string = folder_list[ vid ] + ss.str() + "_gray_" + pix_res.str();
               imwrite( color_save_string, cap_color[ vid ] );
               imwrite( gray_save_string, cap_gray[ vid ] );
            }
         }
         
         imwrite( combined_dir + ss.str() + "_color_" + pix_res.str(), combined_color );
         imwrite( combined_dir + ss.str() + "_gray_" + pix_res.str(), combined_gray );
         
         counter = 0;
         
         cout << endl;
         cout << "Captured Pic: " << num_pic << endl;
         
         num_pic++;
      }
      
      cout << "Time to pic: " << ( delay_s - elapsed ) << endl;
      
      //Waits 30ms and exits loop if user hits key
      if(waitKey(30) >= 0) break;
      
      counter++;
   }
   
   return 0;
}

/////////////////////////
// function: calibrate
// Authors: Brian Fehrman and Scott Logan and Matt Richard
// 
// Purpose: Calibrate cameras based on a given set of images. Can
//          calibrate any number of image sets. Stores the results
//          in the calib folder.
//
//////////////////////////
int calibrate( FileNode fs )
{
   Size board_size;
   string base_dir;
   string img_suf;
   vector< string > folder_names;
   int square_size;
   string calib_dir;
   
   fs[ "board_width" ] >> board_size.width;
   fs[ "board_height" ] >> board_size.height;
   fs[ "base_dir" ] >> base_dir;
   fs[ "img_suf" ] >> img_suf;
   fs[ "square_size" ] >> square_size;
   fs[ "folder_names" ] >> folder_names;
   fs[ "calib_dir" ] >> calib_dir;

   vector<vector<Point3f> > obj_points(1);
   
   base_dir += "/";

   mkdir( calib_dir.c_str(), S_IRWXU|S_IRGRP|S_IXGRP );
   
   calib_dir += "/";
   
   //Approximates position of board corners, needed for calibration
   for( int i = 0; i < board_size.height; ++i )
      for( int j = 0; j < board_size.width; ++j )
         obj_points[ 0 ].push_back(Point3f(float( j*square_size ), float( i*square_size ), 0));

   //Calibrate each image set and save the results
   #pragma omp parallel for
   for( unsigned int img_set = 0; img_set < folder_names.size(); img_set++ )
   {
      vector< vector< Point2f > > all_corners;
      Mat camera_mat = Mat::eye(3, 3, CV_64F );
      Mat dist_coeffs = Mat::zeros( 8, 1, CV_64F );
      vector<Mat> r_vecs, t_vecs;
      vector<Mat> image_set;
      double rms = 0;
      string img_base = base_dir + folder_names[ img_set ] + "/";
      FileStorage calib_save( calib_dir + folder_names[ img_set ] + ".xml", FileStorage::WRITE );
      
      cout << "Calibrating set: " << img_base << endl;
     
      cout << "Finding corners...." << endl;
      
      find_corners( img_base, img_suf, board_size, all_corners, image_set );
      
      if( all_corners.size() > 0 )
      {
         cout << "Succesfully found corners" << endl;
      }
      else
      {
         cout << "Error finding corners" << endl;
         continue;
      }
      
      obj_points.resize( all_corners.size(), obj_points[ 0 ] );
      
      rms = calibrateCamera( obj_points, all_corners, image_set[ 0 ].size(),  camera_mat,
                              dist_coeffs, r_vecs, t_vecs );
                              
      cout << "RMS on Calibration: " << rms << endl;
      
      calib_save << "camera_mat" << camera_mat;
      calib_save << "dist_coeffs" << dist_coeffs;
      calib_save << "r_vecs" << r_vecs;
      calib_save << "t_vecs" << t_vecs;
      calib_save << "rms" << rms;
      
      cout << endl << endl;
      
      calib_save.release();
   }
   
   return 0;
}

/////////////////////////
// function: capture_images
// Authors: Brian Fehrman and Scott Logan and Matt Richard
// Reference: http://docs.opencv.org/doc/tutorials/calib3d/camera_calibration/camera_calibration.html
// 
// Purpose: Find the inner corners in one or more images containg a chessboard
//          and return those corners.
//
//////////////////////////
int find_corners( string img_base, string img_suf, Size board_size,
                     vector< vector< Point2f > > &all_corners, vector<Mat> &image_set )
{   
   int img = 0;
   bool image_found = true;
   string last_image_read;
   
   //For each image in the set, find the corners
   //Continues until an image read is unsuccesful
   do
   {
      Mat image;
      stringstream ss;
      string image_name;
      vector<Point2f> image_corners;
      
      ss << img;
      
      image_name = img_base + ss.str() + img_suf;
      image = imread(  image_name );
      
      if( !image.data )
      {
         if( img > 0 )
         {
            cout << "Final image in set was: " << last_image_read << endl;
         }
         else
         {
            cout << "No images found in set" << endl;
         }
         image_found = false;
      }
      
      else
      {
         bool found = findChessboardCorners( image,
                              board_size, image_corners); //, CV_CALIB_CB_NORMALIZE_IMAGE );
         
         if( !found )
            cout << "Problem finding corners for img: " << img << endl;
         else
         {
            //The cornerSubPix functions improves pixel accuracy
            Mat gray_conv;
            cvtColor( image, gray_conv, CV_BGR2GRAY);
            cornerSubPix( gray_conv, image_corners, Size(11,11),
                        Size(-1,-1), TermCriteria( CV_TERMCRIT_EPS+CV_TERMCRIT_ITER, 30, 0.1 ));
                        
            all_corners.push_back( image_corners );
            image_set.push_back( image );
         }
      }
      
      img++;
      last_image_read = image_name;
      
   }while( image_found );
   
   return 0;
}
